import time

class wm_mpu6050:
    """MPU6050 6轴运动传感器驱动"""
    
    # MPU6050寄存器地址
    MPU6050_ADDR = 0x68  # 默认I2C地址
    
    # 寄存器地址
    REG_SMPLRT_DIV = 0x19      # 采样率分频器
    REG_CONFIG = 0x1A           # 配置寄存器
    REG_GYRO_CONFIG = 0x1B      # 陀螺仪配置
    REG_ACCEL_CONFIG = 0x1C     # 加速度计配置
    REG_FIFO_EN = 0x23          # FIFO使能
    REG_INT_ENABLE = 0x38       # 中断使能
    REG_INT_STATUS = 0x3A       # 中断状态
    REG_ACCEL_XOUT_H = 0x3B     # 加速度计X轴高字节
    REG_ACCEL_XOUT_L = 0x3C     # 加速度计X轴低字节
    REG_ACCEL_YOUT_H = 0x3D     # 加速度计Y轴高字节
    REG_ACCEL_YOUT_L = 0x3E     # 加速度计Y轴低字节
    REG_ACCEL_ZOUT_H = 0x3F     # 加速度计Z轴高字节
    REG_ACCEL_ZOUT_L = 0x40     # 加速度计Z轴低字节
    REG_TEMP_OUT_H = 0x41       # 温度高字节
    REG_TEMP_OUT_L = 0x42       # 温度低字节
    REG_GYRO_XOUT_H = 0x43      # 陀螺仪X轴高字节
    REG_GYRO_XOUT_L = 0x44      # 陀螺仪X轴低字节
    REG_GYRO_YOUT_H = 0x45      # 陀螺仪Y轴高字节
    REG_GYRO_YOUT_L = 0x46      # 陀螺仪Y轴低字节
    REG_GYRO_ZOUT_H = 0x47      # 陀螺仪Z轴高字节
    REG_GYRO_ZOUT_L = 0x48      # 陀螺仪Z轴低字节
    REG_PWR_MGMT_1 = 0x6B       # 电源管理1
    REG_PWR_MGMT_2 = 0x6C       # 电源管理2
    REG_WHO_AM_I = 0x75         # 设备ID
    
    # 陀螺仪量程
    GYRO_RANGE_250 = 0x00       # ±250°/s
    GYRO_RANGE_500 = 0x08       # ±500°/s
    GYRO_RANGE_1000 = 0x10      # ±1000°/s
    GYRO_RANGE_2000 = 0x18      # ±2000°/s
    
    # 加速度计量程
    ACCEL_RANGE_2G = 0x00       # ±2g
    ACCEL_RANGE_4G = 0x08       # ±4g
    ACCEL_RANGE_8G = 0x10       # ±8g
    ACCEL_RANGE_16G = 0x18      # ±16g
    
    def __init__(self, i2c=None, addr=None):
        """
        初始化MPU6050
        
        Args:
            i2c: I2C对象，如果为None则使用全局i2c
            addr: I2C地址，如果为None则使用默认地址0x68
        """
        if i2c is None:
            # 使用全局I2C对象
            try:
                import __main__
                if hasattr(__main__, 'i2c'):
                    self.i2c = __main__.i2c
                else:
                    import builtins
                    if hasattr(builtins, 'i2c'):
                        self.i2c = builtins.i2c
                    else:
                        raise RuntimeError("未找到I2C对象，请先初始化I2C")
            except Exception as e:
                raise RuntimeError(f"未找到I2C对象，请先初始化I2C: {e}")
        else:
            self.i2c = i2c
            
        self.addr = addr or self.MPU6050_ADDR
        
        # 验证设备
        if not self._check_id():
            raise RuntimeError(f"未找到MPU6050芯片 (地址: 0x{self.addr:02X})")
        
        # 初始化配置
        self._init_config()
        
        # 校准参数
        self.accel_offset = [0, 0, 0]
        self.gyro_offset = [0, 0, 0]
        
        print(f"MPU6050初始化成功 (地址: 0x{self.addr:02X})")
    
    def _check_id(self):
        """检查设备ID"""
        try:
            device_id = self._read_reg(self.REG_WHO_AM_I)
            return device_id == 0x68  # MPU6050的设备ID
        except:
            return False
    
    def _init_config(self):
        """初始化配置"""
        # 唤醒设备
        self._write_reg(self.REG_PWR_MGMT_1, 0x00)
        
        # 设置采样率为1kHz
        self._write_reg(self.REG_SMPLRT_DIV, 0x07)
        
        # 设置低通滤波器
        self._write_reg(self.REG_CONFIG, 0x06)
        
        # 设置陀螺仪量程为±2000°/s
        self._write_reg(self.REG_GYRO_CONFIG, self.GYRO_RANGE_2000)
        
        # 设置加速度计量程为±2g
        self._write_reg(self.REG_ACCEL_CONFIG, self.ACCEL_RANGE_2G)
        
        # 设置比例因子
        self.gyro_scale = 16.4      # ±2000°/s量程的比例因子
        self.accel_scale = 16384.0  # ±2g量程的比例因子
    
    def _read_reg(self, reg):
        """读取寄存器"""
        data = self.i2c.readfrom_mem(self.addr, reg, 1)
        return data[0]
    
    def _write_reg(self, reg, value):
        """写入寄存器"""
        self.i2c.writeto_mem(self.addr, reg, bytes([value]))
    
    def _read_word(self, reg):
        """读取16位数据"""
        data = self.i2c.readfrom_mem(self.addr, reg, 2)
        value = (data[0] << 8) | data[1]
        # 处理负数
        if value > 0x7FFF:
            value -= 0x10000
        return value
    
    def get_accel_data(self):
        """
        获取加速度计数据
        
        Returns:
            tuple: (x, y, z) 单位: g
        """
        x = self._read_word(self.REG_ACCEL_XOUT_H) / self.accel_scale
        y = self._read_word(self.REG_ACCEL_YOUT_H) / self.accel_scale
        z = self._read_word(self.REG_ACCEL_ZOUT_H) / self.accel_scale
        
        # 应用校准偏移
        x -= self.accel_offset[0]
        y -= self.accel_offset[1]
        z -= self.accel_offset[2]
        
        return (x, y, z)
    
    def get_gyro_data(self):
        """
        获取陀螺仪数据
        
        Returns:
            tuple: (x, y, z) 单位: °/s
        """
        x = self._read_word(self.REG_GYRO_XOUT_H) / self.gyro_scale
        y = self._read_word(self.REG_GYRO_YOUT_H) / self.gyro_scale
        z = self._read_word(self.REG_GYRO_ZOUT_H) / self.gyro_scale
        
        # 应用校准偏移
        x -= self.gyro_offset[0]
        y -= self.gyro_offset[1]
        z -= self.gyro_offset[2]
        
        return (x, y, z)
    
    def get_temp(self):
        """
        获取温度数据
        
        Returns:
            float: 温度值 (°C)
        """
        temp_raw = self._read_word(self.REG_TEMP_OUT_H)
        temp = temp_raw / 340.0 + 36.53
        return temp
    
    def get_all_data(self):
        """
        获取所有传感器数据
        
        Returns:
            dict: 包含加速度、陀螺仪、温度的数据字典
        """
        accel = self.get_accel_data()
        gyro = self.get_gyro_data()
        temp = self.get_temp()
        
        return {
            'accel': accel,
            'gyro': gyro,
            'temp': temp
        }
    
    def calibrate(self, samples=100):
        """
        校准传感器
        
        Args:
            samples: 校准样本数量，默认100
        """
        print("开始校准MPU6050...")
        print("请保持设备静止不动...")
        
        accel_sum = [0, 0, 0]
        gyro_sum = [0, 0, 0]
        
        for i in range(samples):
            # 读取原始数据
            accel_x = self._read_word(self.REG_ACCEL_XOUT_H) / self.accel_scale
            accel_y = self._read_word(self.REG_ACCEL_YOUT_H) / self.accel_scale
            accel_z = self._read_word(self.REG_ACCEL_ZOUT_H) / self.accel_scale
            
            gyro_x = self._read_word(self.REG_GYRO_XOUT_H) / self.gyro_scale
            gyro_y = self._read_word(self.REG_GYRO_YOUT_H) / self.gyro_scale
            gyro_z = self._read_word(self.REG_GYRO_ZOUT_H) / self.gyro_scale
            
            accel_sum[0] += accel_x
            accel_sum[1] += accel_y
            accel_sum[2] += accel_z - 1.0  # Z轴减去重力加速度
            
            gyro_sum[0] += gyro_x
            gyro_sum[1] += gyro_y
            gyro_sum[2] += gyro_z
            
            import time
            time.sleep(0.01)
        
        # 计算平均值作为偏移
        self.accel_offset = [x / samples for x in accel_sum]
        self.gyro_offset = [x / samples for x in gyro_sum]
        
        print("校准完成!")
        print(f"加速度偏移: {self.accel_offset}")
        print(f"陀螺仪偏移: {self.gyro_offset}")
    

    
    def monitor_data(self, interval=0.1, duration=None):
        """
        持续监控传感器数据
        
        Args:
            interval: 读取间隔时间(秒)，默认0.1秒
            duration: 监控时长(秒)，如果为None则持续监控
        """
        print(f"开始监控MPU6050数据 (间隔: {interval}秒)")
        if duration:
            print(f"监控时长: {duration}秒")
        print("按 Ctrl+C 停止监控")
        print("=" * 60)
        
        count = 0
        start_time = time.time()
        
        try:
            while True:
                if duration and (time.time() - start_time) > duration:
                    break
                    
                count += 1
                data = self.get_all_data()
                
                print(f"[{count:4d}] 时间: {time.time() - start_time:6.1f}s")
                print(f"    加速度 (g): X={data['accel'][0]:6.3f}, Y={data['accel'][1]:6.3f}, Z={data['accel'][2]:6.3f}")
                print(f"    陀螺仪(°/s): X={data['gyro'][0]:6.1f}, Y={data['gyro'][1]:6.1f}, Z={data['gyro'][2]:6.1f}")
                print(f"    温度 (°C): {data['temp']:5.1f}")
                print("-" * 40)
                
                time.sleep(interval)
                
        except KeyboardInterrupt:
            print(f"\n监控已停止，共读取 {count} 次")
        
        print(f"监控完成，总用时: {time.time() - start_time:.1f}秒")
